import os
import sys
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
from PIL import Image

# 定义卷积神经网络模型
class SimpleCNN(nn.Module):
    """
    简单的卷积神经网络类，继承自nn.Module。
    该网络包含两个卷积层和两个全连接层，用于图像处理。
    """
    def __init__(self):
        """
        初始化函数。
        创建卷积层和全连接层的实例。
        """
        # 调用父类的初始化方法
        super(SimpleCNN, self).__init__()
        # 定义第一个卷积层，将1个输入通道转换为32个输出通道，使用3x3的卷积核，步长为1，填充为1
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1)
        # 定义第二个卷积层，将32个输入通道转换为64个输出通道，使用3x3的卷积核，步长为1，填充为1
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)
        # 定义第一个全连接层，输入大小为64*7*7，输出大小为128
        self.fc1 = nn.Linear(64 * 7 * 7, 128)
        # 定义第二个全连接层，输入大小为128，输出大小为10，用于分类
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        """
        定义神经网络的前向传播过程。
        
        参数:
        x: 输入的张量，代表一个批次的图像数据。
        
        返回:
        经过神经网络处理后的输出张量。
        """
        # 应用卷积层1，然后使用ReLU激活函数
        x = F.relu(self.conv1(x))
        # 应用最大池化层1，减少图像的空间尺寸
        x = F.max_pool2d(x, kernel_size=2, stride=2)
        # 应用卷积层2，再次使用ReLU激活函数
        x = F.relu(self.conv2(x))
        # 应用最大池化层2，进一步减少图像的空间尺寸
        x = F.max_pool2d(x, kernel_size=2, stride=2)
        # 将数据扁平化，准备输入到全连接层
        x = x.view(-1, 64 * 7 * 7)
        # 应用全连接层1，并使用ReLU激活函数
        x = F.relu(self.fc1(x))
        # 应用全连接层2，得到最终的输出
        x = self.fc2(x)
        return x


# 数据预处理和加载
def get_data_loaders(batch_size=64):
    """
    获取训练和测试数据加载器。
    
    参数:
    batch_size (int): 每个批次的样本数量，默认为64。
    
    返回:
    tuple: 包含训练加载器和测试加载器的元组。
    """
    # 定义数据预处理步骤：首先转换为张量，然后标准化
    transform = transforms.Compose(
        [transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))]
    )

    # 加载训练数据集
    train_dataset = datasets.MNIST(
        root="../data", train=True, download=True, transform=transform
    )
    # 加载测试数据集
    test_dataset = datasets.MNIST(
        root="../data", train=False, download=True, transform=transform
    )

    # 创建训练数据加载器，设置批次大小和打乱数据顺序
    train_loader = DataLoader(
        dataset=train_dataset, batch_size=batch_size, shuffle=True
    )
    # 创建测试数据加载器，设置批次大小，不打乱数据顺序
    test_loader = DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=False)

    return train_loader, test_loader


# 训练模型
# 参数:
#   model: 待训练的模型
#   train_loader: 训练数据加载器，提供训练所需的图像及其标签
#   criterion: 损失函数，用于计算模型输出与实际标签之间的误差
#   optimizer: 优化器，负责更新模型参数以减小损失
#   epochs: 训练周期数，即完整遍历训练数据集的次数
def train_model(model, train_loader, criterion, optimizer, epochs=5):
    # 将模型设置为训练模式
    model.train()
    # 外层循环控制训练周期
    for epoch in range(epochs):
        # 初始化本次周期的累计损失
        running_loss = 0.0
        # 内层循环按批次处理训练数据
        for images, labels in train_loader:
            # 在每次迭代开始时，清除之前的梯度信息
            optimizer.zero_grad()
            # 计算模型的前向传播输出
            outputs = model(images)
            # 计算本次迭代的损失
            loss = criterion(outputs, labels)
            # 反向传播，计算梯度
            loss.backward()
            # 优化器更新模型参数
            optimizer.step()
            # 累加本次迭代的损失
            running_loss += loss.item()
        # 打印每个训练周期的平均损失
        print(f"Epoch {epoch + 1}/{epochs}, Loss: {running_loss/len(train_loader)}")


# 测试模型
# 参数:
# - model: 待测试的模型
# - test_loader: 测试数据加载器，用于加载测试数据
# 该函数不返回值，它计算并打印模型在测试数据上的准确率
def test_model(model, test_loader):
    # 将模型设置为评估模式
    model.eval()
    # 初始化正确预测的数量和总样本数量
    correct = 0
    total = 0
    # 禁止梯度计算，因为测试阶段不需要反向传播
    with torch.no_grad():
        # 遍历测试数据加载器中的每个样本
        for images, labels in test_loader:
            # 使用模型预测输出
            outputs = model(images)
            # 获取预测结果中概率最高的类别
            _, predicted = torch.max(outputs, 1)
            # 更新总样本数量
            total += labels.size(0)
            # 更新正确预测的数量
            correct += (predicted == labels).sum().item()
    # 计算准确率
    accuracy = correct / total
    # 打印测试准确率
    print(f"Test Accuracy: {accuracy * 100:.2f}%")

def main():
    """
    主函数用于处理模型的训练、测试和预测。
    """
    # 设置训练的批大小、周期数和学习率
    batch_size = 64
    epochs = 5
    learning_rate = 0.001

    # 获取训练和测试的数据加载器
    train_loader, test_loader = get_data_loaders(batch_size)

    # 初始化简单的卷积神经网络模型
    model = SimpleCNN()
    # 定义交叉熵损失函数
    criterion = nn.CrossEntropyLoss()
    # 使用Adam优化器
    optimizer = optim.Adam(model.parameters(), lr=learning_rate)

    # 定义模型保存路径和图片路径
    pth_file = os.path.join(os.getcwd(), "model.pth")

    # 检查模型文件是否存在
    if os.path.exists(pth_file):
        # 如果模型文件存在，删除
        os.remove(pth_file)

    # 训练模型
    print("正在训练模型...")
    train_model(model, train_loader, criterion, optimizer, epochs)

    # 测试模型
    print("正在测试模型...")
    test_model(model, test_loader)

    # 导出模型为 ONNX 格式
    dummy_input = torch.randn(1, 1, 28, 28)  # MNIST 的图像大小为 28x28
    onnx_file_path = "mnist_model.onnx"
    torch.onnx.export(model, dummy_input, onnx_file_path, input_names=["input"], output_names=["output"])
    print(f"模型已导出为 {onnx_file_path}")


if __name__ == "__main__":
    main()
